import datetime

import matplotlib.finance as mpf
import matplotlib.pyplot as plt

import tushare as ts
from matplotlib.pylab import date2num, DateFormatter
import pandas as pd
from LSTM.variableset import EMPath,EMVar
from emutils import stk_utils
from dateutil.parser import parse
import matplotlib.ticker as ticker

import matplotlib
import matplotlib.dates as mdates
matplotlib.use('Agg')



# def draw_candle(data=None, is_save_disk=False, is_show=False):
#     data_list = []
#     for dates, row in data.iterrows():
#         # 将时间转换为数字
#         date_time = datetime.datetime.strptime(row[0], '%Y-%m-%d %H:%M')
#         t = date2num(date_time)
#         open, close, high, low = row[1:5]
#         datas = (t, open, high, low, close)
#         data_list.append(datas)
#     # xfmt = DateFormatter('%Y-%m-%d %H:%M')
#     # 创建子图
#     fig, ax = plt.subplots()
#     fig.subplots_adjust(bottom=0.2)
#     # ax.xaxis.set_major_formatter(xfmt)
#     # 设置X轴刻度为日期时间
#     mpf._candlestick(ax, data_list, width=0.01, colorup='r', colordown='green')
#     # ax.xaxis_date()
#     ax.autoscale_view()
#     plt.setp(plt.gca().get_xticklabels(), rotation=45, horizontalalignment='right')
#     # plt.xticks(rotation=45)
#     # plt.yticks()
#     # plt.title("股票代码：601558两年K线图")
#     # plt.xlabel("时间")
#     # plt.ylabel("股价（元）")
#
#     # plt.grid()
#     plt.show()
# #     # plt.savefig("line.png")  # 保存图



from matplotlib.finance import candlestick2_ohlc

def draw_day_k(df, ax=None, isShow=False, save_path=None,title=None, **kwargs):
    if not isinstance(df, pd.DataFrame):
        raise Exception('数据类型有误')
    if df.empty :
        return
    data_list = []

    for dates, row in df.iterrows():
        # 将时间转换为数字
        date_time = datetime.datetime.strptime(row[EMVar.date], '%Y-%m-%d')
        t = date2num(date_time)
        open, close, high, low = row[EMVar.open],row[EMVar.close],row[EMVar.high],row[EMVar.low]
        datas = (t, open, high, low, close)
        data_list.append(datas)
    if not ax:
        # 创建子图
        fig, ax = plt.subplots()
        fig.subplots_adjust(bottom=0.2)


    # 设置X轴刻度为日期时间
    # xdate = [parse(i) for i in df[EMVar.date]]
    # ax.xaxis.set_major_locator(ticker.MaxNLocator(6))
    # def mydate(x, pos):
    #     try:
    #         return xdate[int(x)]
    #     except IndexError:
    #         return ''
    # ax.xaxis.set_major_formatter(ticker.FuncFormatter(mydate))

    ax.xaxis_date()
    for tick_label in ax.get_xticklabels():
        tick_label.set_rotation(45)

    plt.yticks()
    ax.set_xlabel("Time")
    ax.set_ylabel("Price(RMB)")
    ax.set_title(title) if title else ''
    mpf.candlestick_ohlc(ax, data_list, width=1, colorup='r', colordown='g')
    # candlestick2_ohlc(ax, df[EMVar.open], df[EMVar.high], df[EMVar.low], df[EMVar.close], width=1, colorup='r', colordown='g')
    ma5 = ax.plot(list(df.index), df[EMVar.MA5], color='g', label='ma5')
    ma10 = ax.plot(list(df.index), df[EMVar.MA10], color='r', label='ma10')
    ma20 = ax.plot(list(df.index), df[EMVar.MA20], color='c', label='ma20')
    ma30 = ax.plot(list(df.index), df[EMVar.MA30], color='m', label='ma30')
    # , prop = zhfont
    ax.legend((ma5, ma10, ma20, ma30), (u'MA5', u'MA10', u'MA20', u'MA30'), loc=2)
    ax.legend(loc='upper left')

    if 'buy_date' in kwargs.keys():
        buy_date = kwargs['buy_date']
        if buy_date in df.index:
            tmp = df.ix[buy_date]
            close = float(tmp[EMVar.close])
            ax.axhline(y=close, ls='-.', c='m')
            ax.axvline(x=buy_date, ls='--',c='m')

    if 'sell_date' in kwargs.keys():
        buy_date = kwargs['sell_date']
        if buy_date in df.index:
            tmp = df.ix[buy_date]
            close = float(tmp[EMVar.close])
            ax.axhline(y=close, ls='-.', c='m')
            ax.axvline(x=buy_date, ls='--', c='m')

    if isShow:
        plt.show()
    if save_path:
        plt.savefig(save_path)  # 保存图


def draw_tick_datas(datas=None, ax=None,title=None, save_path=None, is_show=False,hlines=None):
    if isinstance(datas, list):
        i = 0
        if not ax:
            # 创建子图
            fig, ax = plt.subplots()
            fig.subplots_adjust(bottom=0.2)

        # plt.figure(figsize=(8, 4))  # 创建绘图对象
        for data in datas:
            y = data['price'].tolist()
            x = range(i,i + data.shape[0])
            ax.plot(x, y, "b", linewidth=1)  # 在当前绘图对象绘图（X轴，Y轴，蓝色虚线，线宽度）
            ax.axvline(x=x[-1], label="1", color='black', linestyle="--")
            i = i + data.shape[0]
        if hlines and len(hlines) and isinstance(hlines, list):
            for item in hlines:
                if 'style' in item.keys():
                    item_style = item['style']
                else:
                    item_style = '--'
                ax.axhline(y=item['y'], color=item['color'], label="1", linestyle=item_style)
        ax.set_xlabel("Time(s)")  # X轴标签
        ax.set_ylabel("Volt")  # Y轴标签
        ax.set_title(title) if title else '' # 图标题
        if is_show:
            plt.show()  # 显示图
        if save_path:
            plt.savefig(save_path)  # 保存图

def draw_min_data(datas=None, ax=None,title=None, save_path=None, is_show=False,hlines=None):

    if isinstance(datas, list):
        i = 0
        if not ax:
            # 创建子图
            fig, ax = plt.subplots()
            fig.subplots_adjust(bottom=0.2)
        # plt.figure(figsize=(8, 4))  # 创建绘图对象
        for data in datas:
            l_ = list(data.columns)
            if 'open' in l_ and 'close' in l_ and 'high' in l_ and 'low' in l_:
                open = data['open'].tolist()
                high = data['high'].tolist()
                low = data['low'].tolist()
                close = data['close'].tolist()
                x = range(i,i + data.shape[0])
                ax.plot(x, open, "b", linewidth=1)  # 在当前绘图对象绘图（X轴，Y轴，蓝色虚线，线宽度）
                ax.plot(x, high, "r", linewidth=1)
                ax.plot(x, low, "g", linewidth=1)
                ax.plot(x, close, "y", linewidth=1)
                ax.axvline(x=x[-1], label="1", color='black', linestyle="--")
                i = i + data.shape[0]
            else:
                continue
        if hlines and len(hlines) and isinstance(hlines, list):
            for item in hlines:
                if 'style' in item.keys():
                    item_style = item['style']
                else:
                    item_style = '--'
                ax.axhline(y=item['y'], color=item['color'], label="1", linestyle=item_style)
        ax.set_xlabel("Time(s)")  # X轴标签
        ax.set_ylabel("Volt")  # Y轴标签
        ax.set_title(title) if title else '' # 图标题
        if is_show:
            plt.show()  # 显示图
        if save_path:
            plt.savefig(save_path)  # 保存图
    if save_path:
        plt.savefig(save_path)  # 保存图


def draw_total_data(code,operationDate,date_index_list, marketPreoffset=-2, marketLateOffset=5, codePreoffset=-2, codeLateOffset=5,savePath=None,hlines=None):
    date_l, _ = stk_utils.em_get_date_offset(date=operationDate, date_index_list=date_index_list, pre_offset=marketPreoffset,
                                             later_offset=marketLateOffset)
    code_date_l, _ = stk_utils.em_get_date_offset(date=operationDate, date_index_list=date_index_list, pre_offset=codePreoffset,
                                             later_offset=codeLateOffset)
    if date_l and code_date_l:
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2)
        stk_his_data = stk_utils.em_get_h_data(code=code, index=False, start=date_l[0], end=date_l[-1])
        sh_his_data = stk_utils.em_get_h_data(code=EMVar.CODE_SH, index=True, start=date_l[0], end=date_l[-1])
        zx_his_data = stk_utils.em_get_h_data(code=EMVar.CODE_ZX, index=True, start=date_l[0], end=date_l[-1])

        code_min_data = []
        for d in code_date_l:
            print('绘制分钟曲线,日期序列:'+str(code_date_l))
            m = stk_utils.em_resample_tick_data(code=code, date=d)
            code_min_data.append(m)

        draw_day_k(stk_his_data, ax=ax2, title=code + '日线')
        draw_day_k(sh_his_data, ax=ax3, title='上证日线')
        draw_day_k(zx_his_data, ax=ax4, title='中小日线')

        draw_min_data(datas=code_min_data, save_path=savePath,
                                   ax=ax1,
                                   is_show=not savePath,
                                   hlines=hlines)
        plt.close()


def _draw_tick_data(data=None, save_path=None ,is_show=False,**kwargs):

    y = data['price'].tolist()
    x = data.index.tolist()
    plt.figure(figsize=(8, 4))  # 创建绘图对象
    plt.plot(x, y, "b", linewidth=1)  # 在当前绘图对象绘图（X轴，Y轴，蓝色虚线，线宽度）
    plt.axhline(**kwargs)
    plt.xlabel("Time(s)")  # X轴标签
    plt.ylabel("Volt")     # Y轴标签
    plt.title("Line plot")  # 图标题
    if is_show:
        plt.show()  # 显示图
    if save_path:
        plt.savefig(save_path)  # 保存图
    plt.close('all')

# plt.figure(figsize=(12, 12))
# fig, (ax1, ax2) = plt.subplots(1,2)
# fig.subplots_adjust(bottom=0.2,hspace=0.1)
# data = ts.get_tick_data('600001',date='2007-01-09')
# data2 = ts.get_tick_data('600001',date='2007-01-10')
# data3 = ts.get_tick_data('600001',date='2007-01-1')
# draw_tick_datas(datas=[data,data2],is_show=False, ax=ax1)
#
# code = '600848'
# data = ts.get_k_data(code=code, ktype='D', start='2014-10-11', end='2014-12-31', autype='qfq')
# # data = ts.get_hist_data('sh', ktype='5')
# draw_day_k(data, isShow=False, ax=ax2,save_path=EMPath.data_root_path(is_common=True)+'XXX.png')

#
# from sysdata import sys_base_data as sysBaseData
# sfund = sysBaseData.get_self_fund('000856')
# data = sfund.get_index(start='2017-01-26', end='2017-05-26', append_now=False)
# data = data['2017-04-26':]
# draw_day_k(data, isShow=True, save_path=EMPath.data_root_path(is_common=True)+'XXX.png', buy_date='2017-05-12')
